"""
时空马尔可夫链计算
由deepseek生成，尚未经过检验，仅供参考
"""

import geopandas as gpd
import pandas as pd
import numpy as np
from libpysal.weights import Queen
import matplotlib.pyplot as plt
import seaborn as sns
import os

# 1. 参数设置
DATA_DIR = 'path/to/shapefiles'  # Shapefile存放路径
START_YEAR = 2000
END_YEAR = 2010  # 示例年份范围
N_STATES = 3     # 离散状态数（如低、中、高）
OUTPUT_DIR = './results'

os.makedirs(OUTPUT_DIR, exist_ok=True)

# 2. 数据加载与预处理
def load_data():
    """加载所有年份的Shapefile并合并为面板数据"""
    dfs = []
    for year in range(START_YEAR, END_YEAR+1):
        gdf = gpd.read_file(os.path.join(DATA_DIR, f'{year}.shp'))
        gdf['year'] = year
        dfs.append(gdf[['COUNTY_ID', 'year', 'precipitation', 'geometry']])
    
    # 合并并确保几何信息一致
    full_gdf = gpd.GeoDataFrame(pd.concat(dfs, ignore_index=True))
    return full_gdf.sort_values(['COUNTY_ID', 'year'])

full_gdf = load_data()

# 3. 离散化降水量（按分位数划分状态）
def discretize_states(gdf, n_states):
    """为每个年份的降水量数据离散化状态"""
    quantiles = np.linspace(0, 1, n_states+1)
    labels = range(n_states)
    
    gdf['state'] = gdf.groupby('year')['precipitation'].transform(
        lambda x: pd.qcut(x, quantiles, labels=labels, duplicates='drop')
    )
    return gdf

full_gdf = discretize_states(full_gdf, N_STATES)

# 4. 构建空间权重矩阵（使用Queen邻接）
def build_spatial_weights(gdf):
    """构建空间权重矩阵（基于第一年的几何形状）"""
    first_year = gdf[gdf['year'] == START_YEAR]
    w = Queen.from_dataframe(first_year, ids='COUNTY_ID')
    return w

w = build_spatial_weights(full_gdf)

# 5. 计算空间滞后状态
def calculate_spatial_lag(gdf, w):
    """计算每个县每年的空间滞后状态（邻居的众数）"""
    all_years = []
    for year in range(START_YEAR, END_YEAR+1):
        year_gdf = gdf[gdf['year'] == year].set_index('COUNTY_ID')
        states = year_gdf['state'].astype(float)
        
        # 计算空间滞后（邻居状态的众数）
        spatial_lag = []
        for county in year_gdf.index:
            neighbors = w.neighbors[county]
            if neighbors:
                neighbor_states = states.loc[neighbors].mode()
                spatial_lag.append(neighbor_states[0] if not neighbor_states.empty else np.nan)
            else:
                spatial_lag.append(np.nan)
        
        year_gdf['spatial_lag'] = spatial_lag
        all_years.append(year_gdf.reset_index())
    
    return pd.concat(all_years)

full_gdf = calculate_spatial_lag(full_gdf, w)

# 6. 构建空间马尔科夫转移矩阵
def build_spatial_markov_matrix(gdf, n_states):
    """构建空间条件下的转移矩阵"""
    # 准备数据：当前状态、下一年状态、空间滞后状态
    transitions = []
    counties = gdf['COUNTY_ID'].unique()
    for county in counties:
        county_data = gdf[gdf['COUNTY_ID'] == county].sort_values('year')
        for i in range(len(county_data)-1):
            current = county_data.iloc[i]
            next_year = county_data.iloc[i+1]
            transitions.append({
                'current_state': current['state'],
                'next_state': next_year['state'],
                'spatial_lag': current['spatial_lag']
            })
    
    trans_df = pd.DataFrame(transitions).dropna()
    
    # 初始化转移矩阵字典
    markov_matrices = {
        lag: np.zeros((n_states, n_states)) 
        for lag in range(n_states)
    }
    
    # 统计转移次数
    for _, row in trans_df.iterrows():
        current = int(row['current_state'])
        nxt = int(row['next_state'])
        lag = int(row['spatial_lag'])
        
        if not np.isnan(lag) and lag in markov_matrices:
            markov_matrices[lag][current, nxt] += 1
    
    # 转换为概率矩阵
    for lag in markov_matrices:
        row_sums = markov_matrices[lag].sum(axis=1, keepdims=True)
        markov_matrices[lag] = np.divide(
            markov_matrices[lag], row_sums, 
            out=np.zeros_like(markov_matrices[lag]), 
            where=row_sums!=0
        )
    
    return markov_matrices

markov_matrices = build_spatial_markov_matrix(full_gdf, N_STATES)

# 7. 结果输出与可视化
def save_results(matrices, output_dir):
    """保存转移矩阵为CSV和热力图"""
    for lag in matrices:
        df = pd.DataFrame(
            matrices[lag],
            index=[f'State_{i}' for i in range(N_STATES)],
            columns=[f'State_{j}' for j in range(N_STATES)]
        )
        
        # 保存CSV
        df.to_csv(os.path.join(output_dir, f'spatial_lag_{lag}_transition_matrix.csv'))
        
        # 绘制热力图
        plt.figure(figsize=(8,6))
        sns.heatmap(df, annot=True, cmap='YlGnBu', vmin=0, vmax=1)
        plt.title(f'Transition Matrix under Spatial Lag {lag}')
        plt.savefig(os.path.join(output_dir, f'lag_{lag}_heatmap.png'))
        plt.close()

save_results(markov_matrices, OUTPUT_DIR)

print(f"分析结果已保存至：{OUTPUT_DIR}")